from docx import Document
from docx.shared import Pt, RGBColor
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
import os

class DocGenerator:
    def __init__(self):
        self.doc = Document()
    
    def generate_doc(self, original_text: str, ai_analysis: str, output_path: str) -> str:
        """
        生成Word文档
        :param original_text: PDF原文内容
        :param ai_analysis: AI分析结果
        :param output_path: 输出路径
        :return: 生成的文件名
        """
        # 设置文档标题
        title = self.doc.add_heading('PDF文档分析报告', level=0)
        title.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
        
        # 添加原文内容
        self.doc.add_heading('原文内容', level=1)
        original_para = self.doc.add_paragraph(original_text)
        
        # 添加分隔线
        self.doc.add_paragraph('_' * 50)
        
        # 添加AI分析结果
        self.doc.add_heading('AI分析结果', level=1)
        ai_para = self.doc.add_paragraph(ai_analysis)
        
        # 生成唯一的文件名
        filename = f'pdf_analysis_{os.urandom(4).hex()}.docx'
        file_path = os.path.join(output_path, filename)
        
        # 确保输出目录存在
        os.makedirs(output_path, exist_ok=True)
        
        # 保存文档
        self.doc.save(file_path)
        return filename 